'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 3_prompt_router.py
* @Time: 2025/9/15
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_community.utils.math import cosine_similarity
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import OpenAIEmbeddings, ChatOpenAI

dotenv.load_dotenv()

# 1、定义两份不同的模板
physics_template = """你是一位非常聪明的物理教程。
你擅长以简洁易懂的方式回答物理问题。
当你不知道问题的答案时，你会坦率承认自己不知道。

这是一个问题：
{query}"""
math_template = """你是一位非常优秀的数学家。你擅长回答数学问题。
你之所以如此优秀，是因为你能将复杂的问题分解成多个小步骤。
并且回答这些小步骤，然后将它们整合在一起回来更广泛的问题。

这是一个问题：
{query}"""

# 2、创建文本嵌入模型，并执行嵌入
embedding = OpenAIEmbeddings(model="text-embedding-3-small")
prompt_template = [physics_template, math_template]
prompt_embd = embedding.embed_documents(prompt_template)

def prompt_router(input:str)->ChatPromptTemplate:
    """根据传递的query计算返回不同的的提示词模板"""
    query_embedding = embedding.embed_query(input["query"])

    #计算相似性
    similarity = cosine_similarity([query_embedding], prompt_embd)[0]
    most_sm = prompt_template[similarity.argmax()]
    print("使用数学模板" if most_sm == math_template else "使用物理模板")

    #构建提示词模板
    return ChatPromptTemplate.from_template(most_sm)
chain = (
    {"query":RunnablePassthrough()}
    |RunnableLambda(prompt_router)
    |ChatOpenAI(model="gpt-4o-mini")
    |StrOutputParser()
)

print(chain.invoke("什么是非牛顿体"))